{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PXNm5_p_oxMF"
      },
      "source": [
        "# Gemma - Run with Mesop\n",
        "\n",
        "This notebook demonstrates how you can run a Gemma model with [Google Mesop](https://github.com/google/mesop). Mesop is a Python-based UI framework that allows you to rapidly build web apps like demos and internal apps.\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Gemma/[Gemma_2]Using_with_Mesop.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mERVCCsGUPIJ"
      },
      "source": [
        "## Setup\n",
        "\n",
        "### Select the Colab runtime\n",
        "To complete this tutorial, you'll need to have a Colab runtime with sufficient resources to run the Gemma model. In this case, you can use a T4 GPU:\n",
        "\n",
        "1. In the upper-right of the Colab window, select **▾ (Additional connection options)**.\n",
        "2. Select **Change runtime type**.\n",
        "3. Under **Hardware accelerator**, select **L4** or **A100 GPU**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QQ6W7NzRe1VM"
      },
      "source": [
        "### Gemma setup\n",
        "\n",
        "To complete this tutorial, you'll first need to complete the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup). The Gemma setup instructions show you how to do the following:\n",
        "\n",
        "* Get access to Gemma on kaggle.com.\n",
        "* Select a Colab runtime with sufficient resources to run\n",
        "  the Gemma 2B model.\n",
        "* Generate and configure a Kaggle username and API key.\n",
        "\n",
        "After you've completed the Gemma setup, move on to the next section, where you'll set environment variables for your Colab environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_gN-IVRC3dQe"
      },
      "source": [
        "### Set environment variables\n",
        "\n",
        "Set environment variables for `KAGGLE_USERNAME` and `KAGGLE_KEY`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DrBoa_Urw9Vx"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "\n",
        "# Note: `userdata.get` is a Colab API. If you're not using Colab, set the env\n",
        "# vars as appropriate for your system.\n",
        "os.environ[\"KAGGLE_USERNAME\"] = userdata.get(\"KAGGLE_USERNAME\")\n",
        "os.environ[\"KAGGLE_KEY\"] = userdata.get(\"KAGGLE_KEY\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z9oy3QUmXtSd"
      },
      "source": [
        "### Install dependencies\n",
        "\n",
        "You will run KerasNLP to run Gemma. So install Keras and KerasNLP."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UcGLzDeQ8NwN"
      },
      "outputs": [],
      "source": [
        "# Install Keras 3 last. See https://keras.io/getting_started/ for more details.\n",
        "!pip install -q -U keras-nlp\n",
        "!pip install -q -U keras"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pm5cVOFt5YvZ"
      },
      "source": [
        "### Select a backend\n",
        "\n",
        "You will use the JAX backend for this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7rS7ryTs5wjf"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"KERAS_BACKEND\"] = \"jax\"  # Or \"tensorflow\" or \"torch\".\n",
        "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"1.0\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "599765c72722"
      },
      "source": [
        "### Import packages\n",
        "\n",
        "Import Keras and KerasNLP."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f2fa267d75bc"
      },
      "outputs": [],
      "source": [
        "import keras\n",
        "import keras_nlp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xfUlIT24giK8"
      },
      "source": [
        "Enable mixed precision on GPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s79GrIXQf2HS"
      },
      "outputs": [],
      "source": [
        "# Run at half precision.\n",
        "keras.config.set_floatx(\"bfloat16\")\n",
        "keras.mixed_precision.set_global_policy(\"mixed_bfloat16\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZsxDCbLN555T"
      },
      "source": [
        "## Create a model\n",
        "\n",
        "Create the Gemma model using the `from_preset` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yygIK9DEIldp"
      },
      "outputs": [],
      "source": [
        "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma2_instruct_2b_en\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y5EMEQJgnfus"
      },
      "source": [
        "## Install and start Mesop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KK-WfjB11DqO"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting mesop\n",
            "  Downloading mesop-0.9.5-py3-none-any.whl (5.1 MB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/5.1 MB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.5/5.1 MB\u001b[0m \u001b[31m13.9 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━━━━━━━━━━━━━\u001b[0m \u001b[32m3.3/5.1 MB\u001b[0m \u001b[31m51.3 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━━━━━━━━━\u001b[0m \u001b[32m3.9/5.1 MB\u001b[0m \u001b[31m37.3 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m \u001b[32m5.1/5.1 MB\u001b[0m \u001b[31m41.9 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.1/5.1 MB\u001b[0m \u001b[31m33.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: absl-py in /usr/local/lib/python3.10/dist-packages (from mesop) (1.4.0)\n",
            "Collecting deepdiff==6.* (from mesop)\n",
            "  Downloading deepdiff-6.7.1-py3-none-any.whl (76 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.6/76.6 kB\u001b[0m \u001b[31m12.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: flask in /usr/local/lib/python3.10/dist-packages (from mesop) (2.2.5)\n",
            "Requirement already satisfied: markdown in /usr/local/lib/python3.10/dist-packages (from mesop) (3.6)\n",
            "Requirement already satisfied: msgpack in /usr/local/lib/python3.10/dist-packages (from mesop) (1.0.8)\n",
            "Requirement already satisfied: protobuf in /usr/local/lib/python3.10/dist-packages (from mesop) (3.20.3)\n",
            "Requirement already satisfied: pydantic in /usr/local/lib/python3.10/dist-packages (from mesop) (2.8.2)\n",
            "Requirement already satisfied: pygments in /usr/local/lib/python3.10/dist-packages (from mesop) (2.16.1)\n",
            "Collecting python-dotenv (from mesop)\n",
            "  Downloading python_dotenv-1.0.1-py3-none-any.whl (19 kB)\n",
            "Collecting watchdog (from mesop)\n",
            "  Downloading watchdog-4.0.1-py3-none-manylinux2014_x86_64.whl (83 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m83.0/83.0 kB\u001b[0m \u001b[31m14.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting ordered-set<4.2.0,>=4.0.2 (from deepdiff==6.*->mesop)\n",
            "  Downloading ordered_set-4.1.0-py3-none-any.whl (7.6 kB)\n",
            "Requirement already satisfied: Werkzeug>=2.2.2 in /usr/local/lib/python3.10/dist-packages (from flask->mesop) (3.0.3)\n",
            "Requirement already satisfied: Jinja2>=3.0 in /usr/local/lib/python3.10/dist-packages (from flask->mesop) (3.1.4)\n",
            "Requirement already satisfied: itsdangerous>=2.0 in /usr/local/lib/python3.10/dist-packages (from flask->mesop) (2.2.0)\n",
            "Requirement already satisfied: click>=8.0 in /usr/local/lib/python3.10/dist-packages (from flask->mesop) (8.1.7)\n",
            "Requirement already satisfied: annotated-types>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from pydantic->mesop) (0.7.0)\n",
            "Requirement already satisfied: pydantic-core==2.20.1 in /usr/local/lib/python3.10/dist-packages (from pydantic->mesop) (2.20.1)\n",
            "Requirement already satisfied: typing-extensions>=4.6.1 in /usr/local/lib/python3.10/dist-packages (from pydantic->mesop) (4.12.2)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from Jinja2>=3.0->flask->mesop) (2.1.5)\n",
            "Installing collected packages: watchdog, python-dotenv, ordered-set, deepdiff, mesop\n",
            "Successfully installed deepdiff-6.7.1 mesop-0.9.5 ordered-set-4.1.0 python-dotenv-1.0.1 watchdog-4.0.1\n",
            "\n",
            "\u001b[32mRunning server on: http://localhost:32123\u001b[0m\n",
            " * Serving Flask app 'mesop.server.server'\n",
            " * Debug mode: off\n"
          ]
        }
      ],
      "source": [
        "!pip install mesop\n",
        "import mesop as me\n",
        "import mesop.labs as mel\n",
        "\n",
        "me.colab_run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0C2CnGxlnlsK"
      },
      "source": [
        "Load the Mesop UI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D1hVvC5b1KI-"
      },
      "outputs": [
        {
          "data": {
            "application/javascript": "(async (port, path, width, height, cache, element) => {\n    if (!google.colab.kernel.accessAllowed && !cache) {\n      return;\n    }\n    element.appendChild(document.createTextNode(''));\n    const url = await google.colab.kernel.proxyPort(port, {cache});\n    const iframe = document.createElement('iframe');\n    iframe.src = new URL(path, url).toString();\n    iframe.height = height;\n    iframe.width = width;\n    iframe.style.border = 0;\n    iframe.allow = [\n        'accelerometer',\n        'autoplay',\n        'camera',\n        'clipboard-read',\n        'clipboard-write',\n        'gyroscope',\n        'magnetometer',\n        'microphone',\n        'serial',\n        'usb',\n        'xr-spatial-tracking',\n    ].join('; ');\n    element.appendChild(iframe);\n  })(32123, \"/chat\", \"100%\", \"400\", false, window.element)",
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "INFO:werkzeug:\u001b[31m\u001b[1mWARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.\u001b[0m\n",
            " * Running on all addresses (::)\n",
            " * Running on http://[::1]:32123\n",
            " * Running on http://[::1]:32123\n",
            "INFO:werkzeug:\u001b[33mPress CTRL+C to quit\u001b[0m\n"
          ]
        }
      ],
      "source": [
        "@me.page(path=\"/chat\")\n",
        "def chat():\n",
        "    mel.chat(transform)\n",
        "\n",
        "\n",
        "def transform(user_prompt: str, history: list[mel.ChatMessage]) -> str:\n",
        "\n",
        "    # Assemble prompt from chat history\n",
        "    prompt = \"\"\n",
        "    for h in history:\n",
        "        prompt += \"<start_of_turn>{role}\\n{content}<end_of_turn>\\n\".format(\n",
        "            role=h.role, content=h.content\n",
        "        )\n",
        "    prompt += \"<start_of_turn>model\\n\"\n",
        "\n",
        "    result = gemma_lm.generate(prompt)\n",
        "    return result[len(prompt) :]\n",
        "\n",
        "\n",
        "me.colab_show(path=\"/chat\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J7enAHFb1C4i"
      },
      "source": [
        "Now you can chat with the Gemma model in the Mesop UI. You can restart the conversation by running the cell above again."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "[Gemma_2]Using_with_Mesop.ipynb",
      "toc_visible": true
    },
    "google": {
      "image_path": "/site-assets/images/marketing/gemma.png",
      "keywords": [
        "examples",
        "gemma",
        "python",
        "quickstart",
        "text"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
